Switch gradient checkpointing default to use_reentrant=False (PyTorch recommended) #4811
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
5 tasks
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
qgallouedec
requested review from
albertvillanova,
edbeeching,
kashif and
lewtun
edbeeching
approved these changes
Jan 13, 2026
kashif
approved these changes
Jan 13, 2026
Member
There was a problem hiding this comment.
Thanks for the fix! I have a couple of questions below.
Additionally, there are 2 new CI errors that I think we should check if they are caused by this PR and then if they should be fixed as well:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
FAILED tests/experimental/test_nash_md_trainer.py::TestNashMDTrainer::test_training_pre_pefted_model_implicit_ref_with_reward_model - RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
FAILED tests/experimental/test_xpo_trainer.py::TestXPOTrainer::test_training_pre_pefted_model_implicit_ref - RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn…ore setting use_reentrant
… and XPO trainers
Member
Author
Fixed, |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Set
use_reentrant=Falseby defaultThis PR defaults gradient checkpointing to
use_reentrant=Falsein TRL when no explicit value is provided.PyTorch now recommends the non-reentrant checkpointing variant, see https://docs.pytorch.org/docs/stable/checkpoint.html. However, Transformers still defaults to
use_reentrant=Truebecause it was explicitly pinned in the past to silence a PyTorch warning during a transition period, and the default was never updated afterward.Until this is fixed upstream and released (see huggingface/transformers#43203), TRL aligns with the current PyTorch recommendation by setting
use_reentrant=Falseby default, while fully preserving any user-provided value.Fixes
Expected to mark a variable ready only onceThis PR fixes an issue that seems unrelated, but is: #4782
Remove
ScriptArguments.gradient_checkpointing_use_reentrantScriptArguments.gradient_checkpointing_use_reentrantexists but is never used. This is misleading, so this PR removes this argument.